package etree

import (
	"encoding/xml"
	"fmt"
	"html"
	"io"
	"net/http"
	"sort"
	"strings"
)

type (
	Element interface {
		Tag() string

		Text() string
		SetText(text string)

		Tail() string
		SetTail(tail string)

		Length() int
		Item(index int) Element

		Append(elements ...Element)
		Insert(index int, element Element)
		Remove(element Element)

		Get(name string) (string, bool)
		Set(name, value string)

		Keys() []string

		Clear()

		Find(path string) Element
		FindAll(path string) []Element

		WriteHTML(w io.Writer)
		WriteXML(w io.Writer)

		iter(tag string) <-chan Element
	}

	element struct {
		tag      string
		text     string
		tail     string
		children []Element
		attrib   map[string]string
	}

	childless struct {
		text string
		tail string
	}

	comment struct {
		childless
	}

	directive struct {
		childless
	}

	procInst struct {
		childless
	}
)

func NewElement(tag string, attrib map[string]string) Element {
	e := new(element)
	e.tag = tag
	e.children = make([]Element, 0)
	e.attrib = make(map[string]string, len(attrib))
	for k, v := range attrib {
		e.attrib[k] = v
	}
	return e
}

func (self *element) Tag() string {
	return self.tag
}

func (self *element) Text() string {
	return self.text
}

func (self *element) SetText(text string) {
	self.text = text
}

func (self *element) Tail() string {
	return self.tail
}

func (self *element) SetTail(tail string) {
	self.tail = tail
}

func (self *element) Length() int {
	return len(self.children)
}

func (self *element) Item(index int) Element {
	return self.children[index]
}

func (self *element) Append(elements ...Element) {
	self.children = append(self.children, elements...)
}

func (self *element) Insert(index int, element Element) {
	self.children = append(self.children[:index], append([]Element{element}, self.children[index:]...)...)
}

func (self *element) Remove(element Element) {
	for i, v := range self.children {
		if v == element {
			self.children = append(self.children[:i], self.children[i+1:]...)
		}
	}
}

func (self *element) Get(name string) (string, bool) {
	v, ok := self.attrib[name]
	return v, ok
}

func (self *element) Set(name, value string) {
	self.attrib[name] = value
}

func (self *element) Keys() []string {
	ks := make([]string, len(self.attrib))
	i := 0
	for k, _ := range self.attrib {
		ks[i] = k
		i++
	}
	return ks
}

func (self *element) Clear() {
	self.text = ""
	self.tail = ""
	self.children = make([]Element, 0)
	self.attrib = make(map[string]string)
}

func (self *element) Find(path string) Element {
	es := find([]Element{self}, path)
	if len(es) > 0 {
		return es[0]
	}
	return nil
}

func (self *element) FindAll(path string) []Element {
	return find([]Element{self}, path)
}

func (self *element) WriteHTML(w io.Writer) {
	if self.tag == "" {
		fmt.Fprint(w, html.EscapeString(self.text))
		for _, c := range self.children {
			c.WriteHTML(w)
		}
	} else {
		fmt.Fprint(w, "<"+self.tag)
		ks := self.Keys()
		sort.Strings(ks)
		for _, k := range ks {
			fmt.Fprintf(w, ` %s="%s"`, k, self.attrib[k])
		}
		if self.text != "" || self.Length() > 0 {
			fmt.Fprint(w, ">")
			tag := strings.ToLower(self.tag)
			if tag == "script" || tag == "style" {
				fmt.Fprint(w, self.text)
			} else {
				fmt.Fprint(w, html.EscapeString(self.text))
			}
			for _, c := range self.children {
				c.WriteHTML(w)
			}
			fmt.Fprint(w, "</"+tag+">")
		} else {
			fmt.Fprint(w, " />")
		}
	}
	fmt.Fprint(w, html.EscapeString(self.tail))
}

func (self *element) WriteXML(w io.Writer) {
	if self.tag == "" {
		xml.Escape(w, []byte(self.text))
		for _, c := range self.children {
			c.WriteXML(w)
		}
	} else {
		fmt.Fprint(w, "<"+self.tag)
		ks := self.Keys()
		sort.Strings(ks)
		for _, k := range ks {
			fmt.Fprintf(w, ` %s="%s"`, k, self.attrib[k])
		}
		if self.text != "" || self.Length() > 0 {
			fmt.Fprint(w, ">")
			xml.Escape(w, []byte(self.text))
			for _, c := range self.children {
				c.WriteXML(w)
			}
			fmt.Fprint(w, "</"+self.tag+">")
		} else {
			fmt.Fprint(w, " />")
		}
	}
	xml.Escape(w, []byte(self.tail))
}

func iterHelper(ch chan<- Element, e Element, tag string) {
	if e.Tag() != "" && (tag == "*" || e.Tag() == tag) {
		ch <- e
	}
	for i := 0; i < e.Length(); i++ {
		iterHelper(ch, e.Item(i), tag)
	}
}

func (self *element) iter(tag string) <-chan Element {
	ch := make(chan Element)
	go func() {
		iterHelper(ch, self, tag)
		close(ch)
	}()
	return ch
}

func (self *childless) Tag() string {
	return ""
}

func (self *childless) Text() string {
	return self.text
}

func (self *childless) SetText(text string) {
	self.text = text
}

func (self *childless) Tail() string {
	return self.tail
}

func (self *childless) SetTail(tail string) {
	self.tail = tail
}

func (self *childless) Length() int {
	return 0
}

func (self *childless) Item(index int) Element {
	panic(self)
}

func (self *childless) Append(elements ...Element) {
	panic(self)
}

func (self *childless) Insert(index int, element Element) {
	panic(self)
}

func (self *childless) Remove(element Element) {
	panic(self)
}

func (self *childless) Get(name string) (string, bool) {
	panic(self)
}

func (self *childless) Set(name, value string) {
	panic(self)
}

func (self *childless) Keys() []string {
	return make([]string, 0)
}

func (self *childless) Clear() {
	self.text = ""
	self.tail = ""
}

func (self *childless) Find(path string) Element {
	return nil
}

func (self *childless) FindAll(path string) []Element {
	return make([]Element, 0)
}

func (self *childless) iter(tag string) <-chan Element {
	ch := make(chan Element)
	close(ch)
	return ch
}

func NewComment(text string) Element {
	c := new(comment)
	c.text = text
	return c
}

func (self *comment) WriteHTML(w io.Writer) {
	fmt.Fprintf(w, "<!--%s-->", html.EscapeString(self.text))
	fmt.Fprintf(w, html.EscapeString(self.tail))
}

func (self *comment) WriteXML(w io.Writer) {
	fmt.Fprintf(w, "<!--%s-->", self.text)
	xml.Escape(w, []byte(self.tail))
}

func NewDirective(text string) Element {
	d := new(directive)
	d.text = text
	return d
}

func (self *directive) WriteHTML(w io.Writer) {
	fmt.Fprintf(w, "<!%s>", html.EscapeString(self.text))
	fmt.Fprintf(w, html.EscapeString(self.tail))
}

func (self *directive) WriteXML(w io.Writer) {
	fmt.Fprintf(w, "<!%s>", self.text)
	xml.Escape(w, []byte(self.tail))
}

func NewProcInst(target, inst string) Element {
	pi := new(procInst)
	pi.text = target + " " + inst
	return pi
}

func (self *procInst) WriteHTML(w io.Writer) {
	fmt.Fprintf(w, "<?%s?>", html.EscapeString(self.text))
	fmt.Fprintf(w, html.EscapeString(self.tail))
}

func (self *procInst) WriteXML(w io.Writer) {
	fmt.Fprintf(w, "<?%s?>", self.text)
	xml.Escape(w, []byte(self.tail))
}

func ParseString(s string) (Element, error) {
	return Parse(strings.NewReader(s))
}

func ParseURL(url string) (Element, error) {
	resp, err := http.Get(url)
	if err != nil {
		return nil, err
	}
	defer resp.Body.Close()
	return Parse(resp.Body)
}

func Parse(r io.Reader) (Element, error) {
	d := xml.NewDecoder(r)
	root := NewElement("", nil)
	es := []Element{root}

	var parent Element = root
	var t xml.Token
	var err error

	for t, err = d.Token(); t != nil; t, err = d.Token() {
		switch token := t.(type) {
		case xml.StartElement:
			e := NewElement(token.Name.Local, nil)
			for _, v := range token.Attr {
				e.Set(v.Name.Local, v.Value)
			}
			parent.Append(e)
			es = append(es, e)
			parent = e
		case xml.EndElement:
			es = es[:len(es)-1]
			parent = es[len(es)-1]
		case xml.CharData:
			s := string(token)
			if parent != nil {
				if ln := parent.Length(); ln > 0 {
					parent.Item(ln - 1).SetTail(s)
				} else {
					parent.SetText(s)
				}
			}
		case xml.Comment:
			c := NewComment(string(token))
			if parent != nil {
				parent.Append(c)
			}
		case xml.Directive:
			d := NewDirective(string(token))
			if parent != nil {
				parent.Append(d)
			}
		case xml.ProcInst:
			pi := NewProcInst(token.Target, string(token.Inst))
			if parent != nil {
				parent.Append(pi)
			}
		}
	}

	if err != io.EOF {
		return nil, err
	}
	return root, nil
}
